import logging
import os

import paramiko

logging.getLogger("paramiko").setLevel(logging.WARNING)


class SshTool:
    def __init__(self, ssh_ip, ssh_port=22, ssh_user="root", ssh_pwd=""):
        """
        初始化
        :param ssh_ip:   IP
        :param ssh_port:  端口号
        :param ssh_user:  用户名
        :param ssh_pwd:  密码
        """
        self.ip = ssh_ip
        self.port = ssh_port
        self.user = ssh_user
        self.pwd = ssh_pwd
        self.config = None
        # 当前文件目录
        current_file_path = os.path.abspath(__file__)
        current_dir = os.path.dirname(current_file_path)
        # 目标文件位置
        target_path = os.path.dirname(current_dir)
        target_path = os.path.join(target_path, "config")
        private_key_path = os.path.join(target_path, "id_ed25519")
        if not os.path.exists(private_key_path):
            # os.makedirs(target_path)
            logging.error(private_key_path, "不存在！")
        else:
            self.private_key_path = private_key_path
        public_key_path = os.path.join(target_path, "id_ed25519.pub")
        if not os.path.exists(public_key_path):
            # os.makedirs(target_path)
            logging.error(public_key_path, "不存在！")
        else:
            self.public_key_path = public_key_path

    def get_pwd_connection(self):
        """
        使用用户名和密码登陆，仅推荐没有上传key的时候用
        :return:
        """
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(
            hostname=self.ip,
            port=self.port,
            username=self.user,
            password=self.pwd,
            timeout=10,
        )
        return client

    def check_or_upload_ssh_key_pub(self):
        """
        使用 Paramiko 上传 ssh 公钥到服务器的 authorized_keys 文件，避免重复添加
        # chmod 600 ~/.ssh/id_ed25519
        # chmod 644 ~/.ssh/id_ed25519.pub
        """
        conn = None
        try:
            # 建立 SSH 连接
            conn = self.get_pwd_connection()

            # 读取公钥内容
            with open(self.public_key_path, "r") as f:
                public_key = f.read().strip()

            # 在远程服务器上执行命令
            commands = [
                "mkdir -p ~/.ssh",
                "chmod 700 ~/.ssh",
                f'grep -qxF "{public_key}" ~/.ssh/authorized_keys || echo "{public_key}" >> ~/.ssh/authorized_keys',
                "chmod 600 ~/.ssh/authorized_keys",
            ]

            for cmd in commands:
                stdin, stdout, stderr = conn.exec_command(cmd)
                exit_status = stdout.channel.recv_exit_status()
                if exit_status != 0:
                    logging.error(f"Command failed: {cmd}")
                    logging.error(stderr.read().decode())
                    return False

            logging.debug("SSH public key check or uploaded successfully.")
            return True

        except Exception as e:
            logging.error(f"Error uploading SSH public key: {str(e)}")
            return False
        finally:
            if conn:
                conn.close()

    def upload_ssh_key_pub_deprecated(self):
        """
        ssh-copy-id 方式上传ssh公钥，目前有报错，暂时废弃,报错原因未知!
        手动执行可以成功，ssh-copy-id -i /home/lhs/Desktop/test/yunwei/config/id_ed25519.pub -p 22 root@x.x.x.x
        paramiko 执行却报错，pub文件是存在的 /usr/bin/ssh-copy-id: ERROR: failed to open ID file '/home/lhs/Desktop/test/yunwei/config/id_ed25519.pub': No such file
        # chmod 600 ~/.ssh/id_ed25519
        # chmod 644 ~/.ssh/id_ed25519.pub
        """
        command = f"ssh-copy-id -i {self.public_key_path} -p {self.port} {self.user}@{self.ip}"
        conn = self.get_pwd_connection()
        logging.info("----------------------- 命令 ----------------------")
        logging.info(command)
        stdin, stdout, stderr = conn.exec_command(command=command)
        # 输入密码
        stdin.write(self.pwd + "\n")
        stdin.flush()
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        logging.info("----------------------- 结果 ----------------------")
        for output in stdoutput:
            logging.info(output.strip())
        for erroutput in stderroutput:
            logging.error(erroutput.strip())
        logging.debug("Command status: %s" % (stdout.channel.recv_exit_status(),))
        if not stdout.channel.recv_exit_status():
            logging.debug("执行成功。")
            conn.close()
            if not stdoutput:
                stdoutput = True
            return True, stdoutput
        else:
            logging.error("执行失败！")
            for output in stderroutput:
                logging.error("%s" % (output,))
            conn.close()
            return False, stderroutput

    def get_key_connection(self):
        """
        推荐使用key登陆，很多命令不需要再次交互输入密码，稳定
        :return:
        """
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        logging.debug("private_key_path: ", self.private_key_path)
        private_key = paramiko.Ed25519Key.from_private_key_file(self.private_key_path)
        client.connect(
            hostname=self.ip, username=self.user, pkey=private_key, timeout=10
        )
        return client

    def rsync_upload(self, local_dir="~/upload/*", upload_dir="/root/test/"):
        """
        在“服务器”上执行rsync命令，上传文件
        rsync -avP ~/test/* root@ip:/root/test/
        :param local_dir: ~/upload/*
        :param upload_dir: /root/test/
        :return:
        """
        command = f"rsync -avP {local_dir} {self.user}@{self.ip}:{upload_dir}"
        self.run_command(command)

    def rsync_download(self, server_dir="/root/test/*", local_dir="~/download/"):
        """
        在“服务器”上执行rsync命令，下载文件
        rsync -avP root@ip:/root/test/* ~/test/
        :param server_dir: /root/test/*
        :param local_dir: ~/download/
        :return:
        """
        command = f"rsync -avP {self.user}@{self.ip}:{server_dir} {local_dir}"
        self.run_command(command)

    def run_command(self, command="whoami"):
        """执行命令，执行时才建立连接，执行完退出连接。
        :param command:
        :return:如果输入多个命令，返回的是最后命令的状态。
        """
        # 默认使用秘钥方式。
        # 先检查或者上传公钥：如果已上传公钥，不需要密码；如果没有公钥，就需要密码。
        self.check_or_upload_ssh_key_pub()
        conn = self.get_key_connection()
        logging.info("----------------------- 命令 ----------------------")
        logging.info(command)
        stdin, stdout, stderr = conn.exec_command(command=command)
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        for output in stdoutput:
            logging.info(output.strip())
        for erroutput in stderroutput:
            # docker compose 命令结果stderr中，实际有很多正确非错误信息,所以这里按行过滤
            ignored_keywords = ["Container", "Created", "Starting", "Started",
                                "Stopping", "Stopped", "Waiting", "Running", "Healthy", "Removed"]
            if any(keyword in erroutput for keyword in ignored_keywords):
                continue
            logging.error(erroutput.strip())
        logging.debug("Command status: %s" % (stdout.channel.recv_exit_status(),))
        if not stdout.channel.recv_exit_status():
            logging.info("执行成功。")
            logging.info("-------------------------------------------------")
            conn.close()
            if not stdoutput:
                stdoutput = True
            return True, stdoutput
        else:
            logging.error("执行失败！")
            logging.info("-------------------------------------------------")
            for output in stderroutput:
                logging.error("%s" % (output,))
            conn.close()
            return False, stderroutput

    def run_sudo_command(self, command="whoami"):
        """执行sudo命令，执行时才建立连接，执行完退出连接。
        :param command:
        :return:如果输入多个命令，返回的是最后命令的状态。
        """
        conn = self.get_key_connection()
        logging.info("----------------------- 命令 ----------------------")
        logging.info(command)
        stdin, stdout, stderr = conn.exec_command(command=command)
        # 输入密码
        stdin.write(self.pwd + "\n")
        stdin.flush()
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        logging.info("----------------------- 结果 ----------------------")
        for output in stdoutput:
            logging.info(output.strip())
        for erroutput in stderroutput:
            logging.error(erroutput.strip())
        logging.debug("Command status: %s" % (stdout.channel.recv_exit_status(),))
        if not stdout.channel.recv_exit_status():
            logging.debug("执行成功。")
            conn.close()
            if not stdoutput:
                stdoutput = True
            return True, stdoutput
        else:
            logging.error("执行失败！")
            for output in stderroutput:
                logging.error("%s" % (output,))
            conn.close()
            return False, stderroutput

    def run_cmd_quit(self, command="whoami", quit_str="q"):
        """
        最后一个交互式命令，输入退出字符。多个交互式命令，不能同时用
        执行命令，执行时才建立连接，执行完退出连接。
        :param command:
        :param quit_str:  例如 q
        :return:如果输入多个命令，返回的是最后命令的状态。
        """
        conn = self.get_key_connection()
        logging.info("----------------------- 命令 ----------------------")
        logging.info(command)
        stdin, stdout, stderr = conn.exec_command(command=command)
        # 输入退出字符
        stdin.write(quit_str)
        stdin.flush()
        stdoutput = [line for line in stdout]
        stderroutput = [line for line in stderr]
        logging.info("----------------------- 结果 ----------------------")
        for output in stdoutput:
            logging.info(output.strip())
        for erroutput in stderroutput:
            logging.error(erroutput.strip())
        logging.debug("Command status: %s" % (stdout.channel.recv_exit_status(),))
        if not stdout.channel.recv_exit_status():
            logging.debug("执行成功。")
            conn.close()
            if not stdoutput:
                stdoutput = True
            return True, stdoutput
        else:
            logging.error("执行失败！")
            for output in stderroutput:
                logging.error("%s" % (output,))
            conn.close()
            return False, stderroutput

    def reset_history(self):
        cmd = """
        # 重置shell历史记录
        unsetenv("HISTORY");
        unsetenv("HISTFILE");
        unsetenv("HISTSAVE");
        unsetenV("HISTZONE");
        unsetenv("HISTLOG");
        setenv("HISTFILE","/dev/null",1);
        setenv("HISTSIZE”,"0",1);
        setenv("HISTFILESIZE","0",1);
        """
        self.run_command(cmd)


if __name__ == "__main__":
    pass
    # ssh = SshTool(ssh_ip=ip, ssh_user=user)
    # ssh.run_command("id;pwd;")  # ssh.rsync_download()  # ssh.rsync_upload()

    # ssh.run_command('id;ll;')  # ssh.run_command('is;pwd;')  # ssh.run_command('is;pww;')

    # ssh.run_sudo_command('id;pwd;')  # ssh.run_sudo_command('id;ll;')  # ssh.run_sudo_command('is;pwd;')  # ssh.run_sudo_command('is;pww;')

    # ssh.run_cmd_quit('systemctl status docker','q')
